#ifndef _CANARY2_CORO_DATA_PROPOGATOR_H_
#define _CANARY2_CORO_DATA_PROPOGATOR_H_

#include <type_traits>
#include <optional>
#include <mutex>
#include <cstddef>
#include <utility>
#include <atomic>
#include <variant>

#include "normal_awaker.h"

#include "../util/assert.h"

namespace wilson::coro
{
    using unknow_function = void(...);
    struct data_buffer
    {
        explicit data_buffer(unknow_function* dtor) noexcept
            : m_dtor{ dtor } {}

        unknow_function* m_dtor;
        ::std::byte m_data[1];
    };

    void deallocate_data_buffer(data_buffer*);

    template<typename T>
    requires (!::std::is_reference_v<T> 
            && ::std::is_nothrow_move_constructible_v<T>)
    void destroy_data_buffer_and_payload(data_buffer* buffer)
    {
        ::std::destroy_at(reinterpret_cast<T*>(&buffer->m_data));
        ::std::destroy_at(buffer);
        deallocate_data_buffer(buffer);
    }

    template<typename T>
    requires (!::std::is_reference_v<T> 
            && ::std::is_nothrow_move_constructible_v<T>)
    data_buffer* make_data_buffer(T&& val)
    {
        data_buffer* db = reinterpret_cast<data_buffer*>(
            ::operator new(sizeof(data_buffer) + sizeof(T)));
        ::std::construct_at(db, reinterpret_cast<unknow_function*>(destroy_data_buffer_and_payload<T>));
        ::std::construct_at(
            reinterpret_cast<T*>(&db->m_data), 
            ::std::forward<T>(val));
        return db;
    }

    template<typename T>
    requires (!::std::is_reference_v<T> 
            && ::std::is_nothrow_move_constructible_v<T>)
    T move_data(data_buffer* data)
    {
        return ::std::move(*reinterpret_cast<T*>(&data->m_data));
    }

    class data_propagator
    {
    public:
        constexpr data_propagator() noexcept = default;
        ~data_propagator() noexcept
        {
            // is `data_buffer`
            if (m_data.index() == 0)
            {
                auto* datap = get<data_buffer*>(m_data);
                if (datap) reinterpret_cast<void (*)(data_buffer*)>(datap->m_dtor)(datap);
            }
        }

        data_propagator(data_propagator&& other)
            : m_data{ other.m_data }
        {
            other.m_data = nullptr;
        }

        template<typename ValueType>
        requires (!::std::is_reference_v<ValueType> 
                && ::std::is_nothrow_move_constructible_v<ValueType>)
        void set_value(ValueType&& val)
        {
            ::std::lock_guard lk{ m_mutex };
            WILSON_ASSERT_MSG(m_data.index() == 0, "ValueType should NOT void!!!");
            if (auto* datap = get<data_buffer*>(m_data)) [[unlikely]]
                destroy_data_buffer_and_payload<ValueType>(datap);
            m_data = make_data_buffer<ValueType>(::std::forward<ValueType>(val));
            if (m_awaker)
            {
                m_awaker->awake();
            }
        }

        void set_value()
        {
            ::std::lock_guard lk{ m_mutex };
            m_data = true;
            if (m_awaker)
            {
                m_awaker->awake();
            }
        }

        template<typename ValueType>
        requires (!::std::is_reference_v<ValueType> 
                && ::std::is_nothrow_move_constructible_v<ValueType>)
        ValueType retrive_value() noexcept 
        {
            WILSON_ASSERT_MSG(m_data.index() == 0, "ValueType should not void");
            return move_data<ValueType>(get<data_buffer*>(m_data));
        }

        void set_awaker(normal_awaker na)
        {
            ::std::lock_guard lk{ m_mutex };
            switch (m_data.index())
            {
            case 0: 
                if (get<data_buffer*>(m_data) != nullptr) 
                {
                    na.awake(); 
                    return;
                }
                break;

            case 1: 
                if (get<bool>(m_data))
                {
                    na.awake(); 
                    return;
                }
                break;
            }

            m_awaker = ::std::move(na);
        }
        
        bool ready() const 
        {
            ::std::lock_guard lk{ m_mutex };
            switch (m_data.index())
            {
            case 0: return get<0>(m_data) != nullptr;
            case 1: return get<1>(m_data);
            }
            return {};
        }

    private:
        mutable ::std::mutex m_mutex;
        ::std::optional<normal_awaker> m_awaker;
        ::std::variant<data_buffer*, bool> m_data;
    };
}

#endif
